{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import yaml\n",
    "import time\n",
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "os.chdir('../')\n",
    "\n",
    "import mmengine\n",
    "from mot_3d.data_protos import BBox\n",
    "from mot_3d.mot import MOTModel\n",
    "from mot_3d.frame_data import FrameData\n",
    "from mot_3d.update_info_data import UpdateInfoData\n",
    "import mot_3d.tracklet as tracklet\n",
    "from mmcv.ops.diff_iou_rotated import box2corners, oriented_box_intersection_2d\n",
    "from mot_3d.association import associate_dets_to_tracks, associate_unmatched_trks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_frame(data, detection_result, i, label, point_cloud_range):\n",
    "    assert data['sample_idx'] == detection_result['sample_idx'][i]\n",
    "    # points = np.fromfile('./data/CODA/' + data['lidar_path'], dtype=np.float32).reshape([-1, 4])\n",
    "    # mask = (points[:, 0] > point_cloud_range[0]) & (points[:, 0] < point_cloud_range[3]) & \\\n",
    "    #     (points[:, 1] > point_cloud_range[1]) & (points[:, 1] < point_cloud_range[4]) & \\\n",
    "    #     (points[:, 2] > point_cloud_range[2]) & (points[:, 2] < point_cloud_range[5])\n",
    "    # points = points[mask]\n",
    "    ego2global = data['ego2global']\n",
    "    pre_labels = detection_result['pre_labels'][i]\n",
    "    mask = pre_labels == label\n",
    "    pre_bboxes = detection_result['pre_bboxes'][i][mask]\n",
    "    pre_scores = detection_result['pre_scores'][i][mask]\n",
    "    pre_labels = pre_labels[mask]\n",
    "    dets = np.concatenate([pre_bboxes, pre_scores[:, None]], axis=1).tolist()\n",
    "    frame_data = FrameData(dets=dets, ego=ego2global, pc=None, det_types=pre_labels, time_stamp=float(data['timestamp']))\n",
    "    frame_data.dets = [BBox.bbox2world(ego2global, det) for det in frame_data.dets]\n",
    "    \n",
    "    gt_bboxes = []\n",
    "    gt_labels = []\n",
    "    gt_ids = []\n",
    "    for instance in data['instances']:\n",
    "        if instance['bbox_label_3d'] == label:\n",
    "            gt_bboxes.append(BBox.bbox2world(ego2global, BBox.array2bbox(instance['bbox_3d'])))\n",
    "            gt_labels.append(instance['bbox_label_3d'])\n",
    "            gt_ids.append(instance['instance_id'])\n",
    "    \n",
    "    return frame_data, gt_bboxes, gt_labels, gt_ids\n",
    "\n",
    "\n",
    "def iou_rotated_2d(box3d1, box3d2):\n",
    "    box1 = box3d1[..., [0, 1, 3, 4, 6]]  # 2d box\n",
    "    box2 = box3d2[..., [0, 1, 3, 4, 6]]\n",
    "    box1 = box1.unsqueeze(1).repeat(1, box2.shape[0], 1)\n",
    "    box2 = box2.unsqueeze(0).repeat(box1.shape[0], 1, 1)\n",
    "    corners1 = box2corners(box1)\n",
    "    corners2 = box2corners(box2)\n",
    "    intersection, _ = oriented_box_intersection_2d(corners1, corners2)\n",
    "    area1 = box1[..., 2] * box1[..., 3]\n",
    "    area2 = box2[..., 2] * box2[..., 3]\n",
    "    iou = intersection / (area1 + area2 - intersection)\n",
    "    return iou"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_info = mmengine.load('./data/CODA/coda_infos_val.pkl')\n",
    "detection_results = mmengine.load('./data/CODA/detection_results.pkl')\n",
    "config_path = './configs/coda_configs/diou.yaml'\n",
    "save_path = './work_dirs/' + time.strftime('%Y%m%d%H%M%S', time.localtime()) + '/'\n",
    "configs = yaml.load(open(config_path, 'r'), Loader=yaml.Loader)\n",
    "point_cloud_range = [-21.0, -21.0, -2.0, 21.0, 21.0, 6.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing 563 frames\n",
      "Processing scene 10 start from frame 5815\n",
      "Processing scene 10 end at frame 6014 (200 / 563)\n",
      "Processing scene 10 start from frame 9200\n",
      "Processing scene 10 end at frame 9399 (400 / 563)\n",
      "Processing scene 10 start from frame 1987\n",
      "Processing scene 10 end at frame 2149 (563 / 563)\n"
     ]
    }
   ],
   "source": [
    "# infer\n",
    "scene = -1\n",
    "frame = -1\n",
    "all_results = []\n",
    "labels = [0, 1, 2]\n",
    "class_names = ['car', 'pedestrian', 'cyclist']\n",
    "trackers = [None for _ in labels]\n",
    "sequence_results = None\n",
    "\n",
    "print('Processing %d frames' % len(data_info['data_list']))\n",
    "for i, data in enumerate(data_info['data_list']):\n",
    "    if data['scene'] != scene or data['frame'] - 1 != frame:\n",
    "        if sequence_results is not None:\n",
    "            all_results.append(sequence_results)\n",
    "            print('Processing scene %s end at frame %s (%d / %d)' % (scene, frame, i, len(data_info['data_list'])))\n",
    "        scene = data['scene']\n",
    "        frame = data['frame']\n",
    "        print('Processing scene %s start from frame %s' % (scene, frame))\n",
    "        trackers = [MOTModel(configs, class_names[label]) for label in labels]\n",
    "        sequence_results = {'token': data['token'], 'data_list': [[] for _ in labels]}\n",
    "    \n",
    "    scene = data['scene']\n",
    "    frame = data['frame']\n",
    "    for j, label in enumerate(labels):\n",
    "        frame_data, gt_bboxes, gt_labels, gt_ids = load_frame(data, detection_results, i, label, point_cloud_range)\n",
    "        results = trackers[j].frame_mot(frame_data)\n",
    "        result_pred_bboxes = [trk[0] for trk in results]\n",
    "        result_pred_ids = [trk[1] for trk in results]\n",
    "        result_pred_states = [trk[2] for trk in results]\n",
    "        result_labels = [trk[3] for trk in results]\n",
    "\n",
    "        frame_result = {}\n",
    "        frame_result['track_ids'] = result_pred_ids\n",
    "        frame_result['track_bboxes'] = [BBox.bbox2array(bbox) for bbox in result_pred_bboxes]\n",
    "        frame_result['track_states'] = result_pred_states\n",
    "        frame_result['track_labels'] = result_labels\n",
    "        frame_result['gt_bboxes'] = [BBox.bbox2array(bbox) for bbox in gt_bboxes]\n",
    "        frame_result['gt_labels'] = gt_labels\n",
    "        frame_result['gt_ids'] = gt_ids\n",
    "        sequence_results['data_list'][j].append(frame_result)\n",
    "\n",
    "all_results.append(sequence_results)\n",
    "print('Processing scene %s end at frame %s (%d / %d)' % (scene, frame, i+1, len(data_info['data_list'])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing 10_5815\n",
      "Processing 10_9200\n",
      "Processing 10_1987\n",
      "Label 0: MOTA: 0.7819, MT: 0.5000, ML: 0.1000\n",
      "TP: 882, FP: 57, FN: 154, ID Switch: 15\n",
      "Label 1: MOTA: 0.7142, MT: 0.5882, ML: 0.1176\n",
      "TP: 3077, FP: 251, FN: 856, ID Switch: 17\n",
      "Label 2: MOTA: 0.3898, MT: 1.0000, ML: 0.0000\n",
      "TP: 52, FP: 28, FN: 7, ID Switch: 1\n"
     ]
    }
   ],
   "source": [
    "# eval\n",
    "tracks_list = [[] for _ in labels]\n",
    "tracks_num_list = [[] for _ in labels]\n",
    "tp_list = [[] for _ in labels]\n",
    "fp_list = [[] for _ in labels]\n",
    "fn_list = [[] for _ in labels]\n",
    "id_switch_list = [[] for _ in labels]\n",
    "mt_list = [[] for _ in labels]\n",
    "ml_list = [[] for _ in labels]\n",
    "iou_shreshold = [0.7, 0.5, 0.5]\n",
    "\n",
    "for results in all_results:\n",
    "    print('Processing %s' % results['token'])\n",
    "    for i in labels:\n",
    "        class_results = results['data_list'][i]\n",
    "        tracks = {}\n",
    "        tp = 0\n",
    "        fp = 0\n",
    "        fn = 0\n",
    "        mt = 0\n",
    "        ml = 0\n",
    "        id_switch = 0\n",
    "        for frame in class_results:\n",
    "            # 加载数据\n",
    "            gt_bboxes = torch.tensor(np.array(frame['gt_bboxes']), dtype=torch.float).cuda()\n",
    "            gt_ids = frame['gt_ids']\n",
    "            track_bboxes = torch.tensor(np.array(frame['track_bboxes']), dtype=torch.float).cuda()\n",
    "            track_ids = np.array(frame['track_ids'])\n",
    "            track_states = frame['track_states']\n",
    "            valid = np.zeros(len(track_bboxes), dtype=bool)\n",
    "            for k in range(len(track_bboxes)):\n",
    "                tokens = track_states[k].split('_')\n",
    "                if tokens[0] == 'birth':\n",
    "                    valid[k] =  True\n",
    "                elif len(tokens) < 3:\n",
    "                    valid[k] =  False\n",
    "                elif tokens[0] == 'alive' and int(tokens[1]) == 1:\n",
    "                    valid[k] =  True\n",
    "            track_bboxes = track_bboxes[valid]\n",
    "            track_ids = track_ids[valid]\n",
    "\n",
    "            # 匹配，并统计TP、FP、FN\n",
    "            if len(gt_bboxes) > 0 and len(track_bboxes) > 0:\n",
    "                # dist_matrix = torch.cdist(track_bboxes[:, :3], gt_bboxes[:, :3])\n",
    "                iou_matrix = iou_rotated_2d(track_bboxes, gt_bboxes)\n",
    "                matched_ids = np.zeros(len(gt_bboxes), dtype=int) - 1\n",
    "                for k in range(len(track_bboxes)):\n",
    "                    max_iou, max_idx = torch.max(iou_matrix[k], 0)\n",
    "                    if max_iou > iou_shreshold[i]:\n",
    "                        matched_ids[max_idx] = track_ids[k]\n",
    "                        iou_matrix[:, max_idx] = 0\n",
    "                        tp += 1\n",
    "                    else:\n",
    "                        fp += 1\n",
    "                fn += len(gt_bboxes) - sum(matched_ids != -1)\n",
    "\n",
    "                # 构建轨迹匹配序列\n",
    "                for k in range(len(gt_bboxes)):\n",
    "                    if tracks.get(gt_ids[k]) is None:\n",
    "                        tracks[gt_ids[k]] = []\n",
    "                    tracks[gt_ids[k]].append(matched_ids[k])\n",
    "                    \n",
    "            else:\n",
    "                fp += len(track_bboxes)\n",
    "                fn += len(gt_bboxes)\n",
    "\n",
    "        # 计算ID切换、MT、ML\n",
    "        for track in tracks.values():\n",
    "            tracked_num = 0\n",
    "            last_id = -1\n",
    "            for id in track:\n",
    "                if id != -1:\n",
    "                    if last_id != -1 and last_id != id:\n",
    "                        id_switch += 1\n",
    "                    last_id = id\n",
    "                    tracked_num += 1\n",
    "            if tracked_num / len(track) > 0.8:\n",
    "                mt += 1\n",
    "            elif tracked_num / len(track) < 0.2:\n",
    "                ml += 1\n",
    "\n",
    "        tracks_list[i].append(tracks)\n",
    "        tracks_num_list[i].append(len(tracks))\n",
    "        tp_list[i].append(tp)\n",
    "        fp_list[i].append(fp)\n",
    "        fn_list[i].append(fn)\n",
    "        id_switch_list[i].append(id_switch)\n",
    "        mt_list[i].append(mt)\n",
    "        ml_list[i].append(ml)\n",
    "\n",
    "# 计算MOTA、MT、ML\n",
    "for i in labels:\n",
    "    mota = 1 - (sum(fp_list[i]) + sum(fn_list[i]) + sum(id_switch_list[i])) / (sum(tp_list[i]) + sum(fn_list[i]))\n",
    "    mt = sum(mt_list[i]) / sum(tracks_num_list[i])\n",
    "    ml = sum(ml_list[i]) / sum(tracks_num_list[i])\n",
    "    print('Label %d: MOTA: %.4f, MT: %.4f, ML: %.4f' % (i, mota, mt, ml))\n",
    "    print('TP: %d, FP: %d, FN: %d, ID Switch: %d' % (sum(tp_list[i]), sum(fp_list[i]), sum(fn_list[i]), sum(id_switch_list[i])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmengine.dump(all_results, './data/CODA/coda_track_result.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_time1 = 0\n",
    "infer_time2 = 0\n",
    "infer_time3 = 0\n",
    "infer_time4 = 0\n",
    "infer_time5 = 0\n",
    "for scene in track_info:\n",
    "    tracker = MOTModel(configs)\n",
    "    scene_results = []\n",
    "    for frame_id, data in enumerate(scene['data_list']):\n",
    "        time0 = time.time()\n",
    "        input_data, gt_bboxes = load_frame(data, point_cloud_range)\n",
    "\n",
    "        time1 = time.time()\n",
    "        tracker.frame_count += 1\n",
    "        # initialize the time stamp on frame 0\n",
    "        if tracker.time_stamp is None:\n",
    "            tracker.time_stamp = input_data.time_stamp\n",
    "    \n",
    "        # filter out low-score detections\n",
    "        dets = input_data.dets\n",
    "        det_indexes = [i for i, det in enumerate(dets) if det.s >= tracker.score_threshold]\n",
    "        dets = [dets[i] for i in det_indexes]\n",
    "\n",
    "        # prediction and association\n",
    "        trk_preds = list()\n",
    "        for trk in tracker.trackers:\n",
    "            trk_preds.append(trk.predict(input_data.time_stamp))\n",
    "        matched, unmatched_dets, unmatched_trks = associate_dets_to_tracks(dets, trk_preds, \n",
    "            tracker.match_type, tracker.asso, tracker.asso_thres)\n",
    "        for k in range(len(matched)):\n",
    "            matched[k][0] = det_indexes[matched[k][0]]\n",
    "        for k in range(len(unmatched_dets)):\n",
    "            unmatched_dets[k] = det_indexes[unmatched_dets[k]]\n",
    "        \n",
    "        time2 = time.time()\n",
    "        # association in second stage\n",
    "        dets = input_data.dets\n",
    "        det_indexes = [i for i, det in enumerate(dets) if det.s >= tracker.score_threshold_second_stage]\n",
    "        dets = [dets[i] for i in det_indexes]\n",
    "        unmatched_trk_preds = [tracker.trackers[t].get_state() for t in unmatched_trks]\n",
    "        update_modes = associate_unmatched_trks(dets, unmatched_trk_preds, tracker.asso, tracker.asso_thres_second_stage)\n",
    "\n",
    "        time3 = time.time()\n",
    "        # update the matched tracks\n",
    "        for i in range(len(matched)):\n",
    "            d = matched[i][0]\n",
    "            trk = tracker.trackers[matched[i][1]]\n",
    "            update_info = UpdateInfoData(mode=1, bbox=input_data.dets[d], ego=input_data.ego, \n",
    "                        frame_index=tracker.frame_count, pc=input_data.pc, dets=input_data.dets)\n",
    "            trk.update(update_info)\n",
    "        for i in range(len(unmatched_trks)):\n",
    "            trk = tracker.trackers[unmatched_trks[i]]\n",
    "            update_info = UpdateInfoData(mode=update_modes[i], bbox=unmatched_trk_preds[i], ego=input_data.ego, \n",
    "                        frame_index=tracker.frame_count, pc=input_data.pc, dets=input_data.dets)\n",
    "            trk.update(update_info)\n",
    "\n",
    "        time4 = time.time()\n",
    "        # create new tracks for unmatched detections\n",
    "        for index in unmatched_dets:\n",
    "            track = tracklet.Tracklet(tracker.configs, tracker.count, input_data.dets[index], input_data.det_types[index], \n",
    "                tracker.frame_count, time_stamp=input_data.time_stamp)\n",
    "            tracker.trackers.append(track)\n",
    "            tracker.count += 1\n",
    "        \n",
    "        # remove dead tracks\n",
    "        track_num = len(tracker.trackers)\n",
    "        for index, trk in enumerate(reversed(tracker.trackers)):\n",
    "            if trk.death(tracker.frame_count):\n",
    "                tracker.trackers.pop(track_num - 1 - index)\n",
    "        \n",
    "        # output the results\n",
    "        results = list()\n",
    "        for trk in tracker.trackers:\n",
    "            state_string = trk.state_string(tracker.frame_count)\n",
    "            results.append((trk.get_state(), trk.id, state_string, trk.det_type))\n",
    "        \n",
    "        # wrap up and update the information about the mot trackers\n",
    "        tracker.time_stamp = input_data.time_stamp\n",
    "        for trk in tracker.trackers:\n",
    "            trk.sync_time_stamp(tracker.time_stamp)\n",
    "\n",
    "        result_pred_bboxes = [trk[0] for trk in results]\n",
    "        result_pred_ids = [trk[1] for trk in results]\n",
    "        result_pred_states = [trk[2] for trk in results]\n",
    "        result_types = [trk[3] for trk in results]\n",
    "\n",
    "        frame_result = {}\n",
    "        frame_result['track_ids'] = result_pred_ids\n",
    "        frame_result['track_bboxes'] = [BBox.bbox2array(bbox) for bbox in result_pred_bboxes]\n",
    "        frame_result['track_states'] = result_pred_states\n",
    "        scene_results.append(frame_result)\n",
    "\n",
    "        time5 = time.time()\n",
    "        infer_time1 += time1 - time0\n",
    "        infer_time2 += time2 - time1\n",
    "        infer_time3 += time3 - time2\n",
    "        infer_time4 += time4 - time3\n",
    "        infer_time5 += time5 - time4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_time1, infer_time2, infer_time3, infer_time4, infer_time5"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
